from interface import Runner
import redis
from time import sleep
import mss
import pandas as pd
import os
from os import path
from paddleocr import PaddleOCR
from pynput import keyboard

# from ghub import click_key
# import pydirectinput as pyinput

class Sector(Runner):
    def __init__(self, redis_pool: redis.ConnectionPool):
        super().__init__(redis_pool)
        self.finish = False
        
    def producter(self):
        # 获取代码列表
        stock_codes = os.listdir(r'E:\new_tdx\T0002\5-export')
        self.code_list = []
        for code in stock_codes:
            self.code_list.append(code[3:9])
        # 实例化截图变量
        self.m = mss.mss()
        self.rect = (400,1455,530,1515)
        # 计数
        pi = 1
        ctrl = keyboard.Controller()
        for code in self.code_list:
            # 进入股票详情
            for cha in code:
                ctrl.press(cha)
                ctrl.release(cha)
                sleep(0.1)
            ctrl.press(keyboard.Key.enter)
            ctrl.release(keyboard.Key.enter)
            sleep(0.8)
            # 截图保存
            img = self.m.grab(self.rect)
            mss.tools.to_png(img.rgb, img.size, 8, f'shot_image/{code}.jpg')
            pi += 1
            if pi%50 == 0:
                print(f'product : {pi}')
        self.finish = True
            

    def consumer(self):
        # 实例化ocr
        self.ocr = PaddleOCR(use_angle_cls=False, lang='en', show_log='INFO')
        # 计数
        ci = 1
        while(True):
            images = os.listdir(r'./shot_image')
            if len(images) == 0:
                sleep(1)
                continue
            with open('band.txt',mode='a',encoding='utf-8') as f:
                for imgname in images:
                    result = self.ocr.ocr(f'shot_image/{imgname}', cls=False)
                    f.write(imgname.split('.')[0]+',')
                    for res in result:
                        f.write(res[1][0]+'\n')
                    os.remove(f'shot_image/{imgname}')
                    ci += 1
                    if ci%50 == 0:
                        print(f'consum : {ci}')
            sleep(0.5)
            if self.finish:
                sleep(3)
                images = os.listdir(r'./shot_image')
                if len(images) == 0:
                    break

redis_pool = redis.ConnectionPool(host='127.0.0.1', port=6379, max_connections=5, db=3, decode_responses=True)
test = Sector(redis_pool)
sleep(2)
test.run()
